import argparse
import os
import torch

def load_sd(path):
    sd = torch.load(path, map_location="cpu")
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    return sd

def average_state_dicts(state_dicts):
    keys = set()
    for sd in state_dicts:
        keys.update(sd.keys())
    merged = {}
    for k in keys:
        tensors = []
        for sd in state_dicts:
            v = sd.get(k, None)
            if isinstance(v, torch.Tensor):
                tensors.append(v)
        if len(tensors) > 0 and all(t.shape == tensors[0].shape for t in tensors):
            dtype = tensors[0].dtype
            acc = torch.zeros_like(tensors[0], dtype=dtype)
            for t in tensors:
                acc = acc + t.to(dtype)
            merged[k] = acc / len(tensors)
        else:
            for sd in state_dicts:
                if k in sd:
                    merged[k] = sd[k]
                    break
    return merged

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--models", type=str, required=True)
    parser.add_argument("--output", type=str, required=True)
    args = parser.parse_args()

    paths = [p for p in args.models.split(",") if p]
    sds = [load_sd(p) for p in paths]
    merged = average_state_dicts(sds)

    outdir = os.path.dirname(args.output)
    if outdir:
        os.makedirs(outdir, exist_ok=True)
    torch.save(merged, args.output)
    print(f"Saved averaged state dict to {args.output}")

if __name__ == "__main__":
    main()

